'''
Ant Colony Optimization Algorithm for TSP
'''
from pre import *
import numpy as np
import numpy.random as rand
import random
import matplotlib.pyplot as plt
import time
import sys
import math as m
import xlrd
import itertools as itt
from matplotlib import cm

'''
consts
'''

n_cities =12
n_ants = 20
cycle_max = 300

g_dist = []
alpha = 1  # importance of tao on decison
beta = 1  # importance of expectation(heuristic) on decison
Q = 20
rho = 0.5
g_pheromones = [[0 for col in range(n_cities)]for row in range(
    n_cities)]  # global tao --pheromone

p_exploit = 0.9

xx = []
yy = []
def heur(i, j):
    # heuristic
    if g_dist[i][j] == 0:
        pass
        p = 9
        p = 10
    return (1/g_dist[i][j])


def roulette(probs):
    total = 0
    cdf = []
    accu = 0
    for p in probs:
        total = total + p
    for i in range(len(probs)):
        accu = accu + probs[i] / total
        cdf.append(accu)
    r = random.random()
    i = 0
    # print("CDF")
    # print(cdf)
    for c in cdf:
        if r < c:
            return i
        i = i + 1


'''
class for ant
'''


class Ant:
    # an ant
    def __init__(self):
        self.pos = random.randint(0, n_cities - 1)  # start position
        self.path = [self.pos]
        self.allow = list(range(n_cities))  # cities remained to be visited
        # remove the
        self.allow.remove(self.pos)
        self.distance = 0

    def proceed(self, exploit=False):
        sum_h = 0
        sum_tao = 0
        probs = []
        best_h_idx = 1e9
        best_h = -1
        allow_idx = 0
        if len(self.allow) == 0:
            self.distance = self.distance + \
                g_dist[self.path[n_cities - 1]][self.path[0]]
            return
        if len(self.allow) == 1:
            # print("FINISH")
            allow_idx = 0
        else:
            if exploit == False:
                # whether use the technology of ASM
                # always use the traditional way
                for i in self.allow:
                    # every city it can go to
                    tao = (g_pheromones[self.pos][i])**alpha
                    h = heur(self.pos, i)**beta
                    sum_h = sum_h + h
                    sum_tao = sum_tao + tao
                for i in range(len(self.allow)):
                    probs.append((h+tao)/sum_h+sum_tao)
                # based on pheromones and heuristic
                # index in allow_succ
                allow_idx = roulette(probs)
            else:
                if random.random() < p_exploit:
                    p = []
                    for i in self.allow:
                        p.append(g_pheromones[self.pos][i]
                                 * (heru(self.pos, i)**beta))
                    allow_idx = 0
                    max_h = -1
                    for i in range(len(p)):
                        if p[i] > max_h:
                            # exploit
                            allow_idx = i
                            max_h = p[i]
                else:
                    for i in self.allow:
                        # every city it can go to
                        tao = (g_pheromones[self.pos][i])**alpha
                        h = heur(self.pos, i)**beta
                        sum_h = sum_h + h
                        sum_tao = sum_tao + tao
                    for i in range(len(self.allow)):
                        probs.append((h+tao)/sum_h+sum_tao)
                        # based on pheromones and heuristic
                        # index in allow_succ
                    allow_idx = roulette(probs)
        # print("ALLOW")
        # print(self.allow)
        self.pos = self.allow[allow_idx]
        self.path.append(self.allow[allow_idx])
        self.allow.remove(self.allow[allow_idx])
        # print("POS")
        # print(self.pos)
        # print("INDEX")
        # print(idx)
        self.distance = self.distance + g_dist[self.pos][allow_idx]
        # print("DIST")
        # print(self.distance)
        # print("PATH")
        # print(self.path)

    def release_pheromone(self):
          # finish its travel
          # leave pheromone
        for i in range(n_cities-1):
            idxa = self.path[i]
            idxb = self.path[i + 1]
            g_pheromones[idxa][idxb] = g_pheromones[idxa][idxb] + \
                Q / self.distance
            g_pheromones[idxb][idxa] = g_pheromones[idxa][idxb]
        idxa = self.path[0]
        idxb = self.path[n_cities-1]
        g_pheromones[idxa][idxb] = g_pheromones[idxa][idxb] + \
            Q / self.distance
        g_pheromones[idxb][idxa] = g_pheromones[idxa][idxb]


'''
functions
'''

# def init_ant_pheromone():
#     # globally init


def g_evaporate():
    # globally change the pheromones
    for i in range(n_cities):
        for j in range(n_cities):
            g_pheromones[i][j] = g_pheromones[i][j] * (1 - rho)


'''
functions for visualization
'''


def init_diagram(x, y, plt):

    x = x[0:n_cities]
    y = y[0:n_cities]

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)

    ax.scatter(x, y)
    ax.set_xlabel("Coord X")
    ax.set_ylabel("Coord Y")

    for i in range(n_cities):
        ax.annotate(str(i), (x[i], y[i]), color='black')

    std_text = ax.texts
    fig.canvas.draw()

    ax_bg = fig.canvas.copy_from_bbox(ax.bbox)

    plt.ion()
    return fig, ax, ax_bg, std_text


def display(fig, ax, ax_bg, std_text, best_ever, best_distance):

    # print(ind)
    # plot accroding to the permulation
    ax.lines = []
    # ax.texts = []
    fig.canvas.restore_region(ax_bg)

    # best_ever
    xs = []
    ys = []
    best_ever = best_ever.path
    for i in range(n_cities):
        ida = best_ever[i]
        # idb = ind[i+1]
        xs.append(x[ida])
        ys.append(y[ida])

    # pheromone
    fmax = max(max(row) for row in g_pheromones)
    for i in range(n_cities):
        for j in range(i + 1, n_cities):
            f = g_pheromones[i][j]
            p, = ax.plot((x[i], x[j]), (y[i], y[j]), color=cm.YlGn(f / fmax))

    xs.append(x[best_ever[0]])
    ys.append(y[best_ever[0]])
    p, = ax.plot(xs, ys, color='r')
    ax.draw_artist(p)
    s = str(best_ever)

    title = ax.set_title(s+"\nSum of dist : "+str(best_distance), wrap=True)
    ax.draw_artist(title)

    for at in std_text:
        ax.draw_artist(at)

    fig.canvas.blit(ax.bbox)


'''
mainn funtion
'''


def aco_tsp(x, y, acs=False):

    # main function of aco
    best_ant_ever = []
    best_distance = 1e9
    cycle = 0
    ants = []
    init_t = time.process_time()
    frame_count = 0

    fig, ax, ax_bg, std_text = init_diagram(x, y, plt)

    for n_cycle in range(cycle_max):
        ants = []
        # init
        for i in range(n_ants):
            a = Ant()
            ants.append(a)

        # evaporate
        g_evaporate()
        for city in range(n_cities):
            for ant in ants:
                # proceed
                ant.proceed()
        ants = sorted(ants, key=lambda ant: ant.distance)
        # whether to take the strategy of Ant Colony System
        if acs == True:
            # select the best
            ants[0].release_pheromone()
        else:
            for ant in ants:
                ant.release_pheromone()
        if ants[0].distance < best_distance:
            best_ant_ever = ants[0]
            best_distance = ants[0].distance

        display(fig, ax, ax_bg, std_text, best_ant_ever, best_distance)
        plt.pause(1e-4)
        frame_count=frame_count + 1
        print("ITER", n_cycle,"BEST",best_distance)
        xx.append(n_cycle)
        yy.append(best_distance)

    print("**************GA DONE**************")
    print("TIME")
    print(time.process_time() - init_t)
    print("FRAME")
    print(frame_count)
    return best_ant_ever


'''
test
'''

g_dist = get_dist_matrix(
    "E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/dist.txt")
xy = np.loadtxt(
    'E:/codes/courses/ArtificialIntelligence/GA/tsp_ga/raw_xy.txt')
x = xy[..., 0][0:n_cities]
y = xy[..., 1][0:n_cities]
# plt.ion()
# plt.show()

with open("E:/codes/courses/ArtificialIntelligence/ACA/out/out.txt",'w+',encoding= 'utf-8') as file:
    sys.stdout = file
    elite = aco_tsp(x, y, True)
    print(elite.distance)
# plt.ioff()
# plt.cla()
# plt.plot(xx, yy)
# plt.show()
# print(g_pheromones)
